'''
Pika extensions, including:
  - pika credentials helper
  - simple connection pool implementation
  - decorator that decorates iterator, it adds a connection from a pool to args
  - blocking connection, suited for AWSGI (evserver)
'''
import errno
import heapq
import pika
import pika.connection
import pika.exceptions
import socket
import time
import functools
import re
import Queue


amqp_re = re.compile(r'^amqp://'\
                '((?P<username>[^:]*)[:](?P<password>[^@]*)[@])?' \
                '((?P<host>[^/:]*)([:](?P<port>[^/]*))?)' \
                '(?P<virtual_host>/[^/]*)/?' \
                '([?](?P<query>[^/]*))?$')

def pika_credentials_from_url(amqp_url = None):
    if not amqp_url:
        amqp_url = 'amqp://127.0.0.1/'
    params = amqp_re.match(amqp_url).groupdict()

    credentials = pika.PlainCredentials(
                        params['username'] or 'guest',
                        params['password'] or 'guest',
                    )
    conn_params = pika.ConnectionParameters(
                        params['host'] or '127.0.0.1',
                        port = int(params['port'] or '5672'),
                        virtual_host = params['virtual_host'] or '/',
                        credentials = credentials,
                    )
    return conn_params


class ConnectionPool:
    def __init__(self, amqp_url):
        self.queue = Queue.Queue()
        self.amqp_url = amqp_url
        self.pika_credentials = pika_credentials_from_url(amqp_url)

    def get(self):
        try:
            return self.queue.get(False)
        except Queue.Empty:
            conn = AwsgiConnection(self.pika_credentials)
            ch = conn.channel()
            return conn

    def put(self, conn):
        self.queue.put(conn, False)


def connection_from_pool(connection_pool):
    def decorator(iterator):
        @functools.wraps(iterator)
        def wrapper(*args, **kwargs):
            conn = connection_pool.get()
            try:
                iterable = iterator(conn, *args, **kwargs)
                try:
                    while True:
                        yield iterable.next()
                except GeneratorExit:
                    iterable.close()
            except pika.exceptions.ConnectionClosed:
                conn.close()
                print "conn died"
            else:
                if conn.is_alive():
                    connection_pool.put( conn )
                else:
                    conn.close()
                    print "conn returned dead"
        return wrapper
    return decorator


class AwsgiConnection(pika.connection.Connection):
    def __init__(self, *args, **kwargs):
        self.timer_heap = []
        pika.connection.Connection.__init__(self, *args, **kwargs)

    def delayed_call(self, delay_sec, callback):
        self.add_oneshot_timer_rel(delay_sec, callback)

    def add_oneshot_timer_rel(self, firing_delay, callback):
        self.add_oneshot_timer_abs(time.time() + firing_delay, callback)

    def add_oneshot_timer_abs(self, firing_time, callback):
        heapq.heappush(self.timer_heap, (firing_time, callback))

    def _run_timers_internal(self):
        cutoff = time.time()
        while self.timer_heap and self.timer_heap[0][0] < cutoff:
            try:
                heapq.heappop(self.timer_heap)[1]()
            except:
                self.log_timer_error(sys.exc_info())
            cutoff = time.time()
        return cutoff

    def next_event_timeout(self):
        cutoff = self._run_timers_internal()
        if self.timer_heap:
            timeout = self.timer_heap[0][0] - cutoff
        else:
            timeout = 300.0 # default timeout
        return timeout


    def connect(self, host, port):
        self.socket = socket.socket()
        self.socket.connect((host, port))
        self.socket.setblocking(False)
        self.on_connected()

    def disconnect_transport(self):
        self.socket.close()
        self.on_disconnected()


    def awsgi_yield(self, request, timeout=Ellipsis):
        ''' Return parameter for Async WSGI 'yield' in iterator '''
        min_timeout = min(timeout, self.next_event_timeout())
        try:
            self._do_send()
        except socket.error, (exn_errno, exn_msg):
            if exn_errno == errno.EAGAIN: # more to be written
                return request.environ['x-wsgiorg.fdevent.writable'](self.socket, min_timeout)
            else:
                raise

        r = self._do_recv()
        if r:
            return ''
        return request.environ['x-wsgiorg.fdevent.readable'](self.socket, min_timeout)

    def drain_events(self, read=True):
        ''' Block and do the send and recv '''
        self.socket.setblocking(True)
        self._do_send()
        if read:
            self._do_recv()
        self.socket.setblocking(False)

    def _do_send(self):
        try:
            while self.outbound_buffer:
                fragment = self.outbound_buffer.read()
                r = self.socket.send(fragment)
                self.outbound_buffer.consume(r)
        except socket.error, (exn_errno, exn_msg):
            if exn_errno in (errno.EPIPE, errno.ECONNRESET):
                self.disconnect_transport()
                raise pika.exceptions.ConnectionClosed()
            if exn_errno == errno.EINTR: # retry
                return self._do_send()
            raise

    def _do_recv(self):
        try:
            buf = self.socket.recv(self.suggested_buffer_size())
        except socket.error, (exn_errno, exn_msg):
            if exn_errno == errno.EINTR: # retry
                return self._do_recv()
            if exn_errno == errno.EAGAIN:
                return 0
            raise
        else:
            if not buf:
                self.disconnect_transport()
                raise pika.exceptions.ConnectionClosed()
            self.on_data_available(buf)
            return len(buf)
